/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "multimodel.h"
#include "binomial.h"

#include <errno.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

#include "normalize.c"

#define MODEL_VERSION (1ULL << 24)

typedef struct _Header Header;
struct _Header 
{
  uint64_t      version;
  uint64_t      n_items;
  uint32_t      n_raw_labels;
  uint32_t      max_raw_labels;
  uint32_t      n_workers;
  uint32_t      rank;
};

typedef struct _MultiModelImpl MultiModelImpl;
struct _MultiModelImpl 
{
  int                                   fd;
  uint8_t*                              base;
  ParameterizedNormalDistribution       priorloga;
  NormalDistribution                    priorlogbeta;
  NormalDistribution                    priorc;
  NormalDistribution                    priord;
  float*                                logpriorz;
};

static unsigned int
n_labels (const Header* h)
{
  unsigned int rv = 0;

  for (unsigned int m = 0; m <= h->max_raw_labels; ++m)
    {
      rv += binomial (h->n_raw_labels, m);
    }

  return rv;
}

static float*
loga_pointer (uint8_t* base)
{
  Header* h = (Header*) base;

  return (float*) (
         base 
       + sizeof (*h) 
       + (h->n_raw_labels * sizeof (float))                      /* logpriorz */
       + (n_labels (h) * h->rank * sizeof (float))               /* c */
       + (n_labels (h) * h->rank * sizeof (float))               /* d */
       + (h->rank * sizeof (float))                              /* gamma */
    );
}

static float*
c_pointer (uint8_t* base)
{
  Header* h = (Header*) base;

  return (float*) (
         base 
       + sizeof (*h) 
       + (h->n_raw_labels * sizeof (float))                      /* logpriorz */
    );
}

static float*
d_pointer (uint8_t* base)
{
  Header* h = (Header*) base;

  return (float*) (
         base 
       + sizeof (*h) 
       + (h->n_raw_labels * sizeof (float))                      /* logpriorz */
       + (n_labels (h) * h->rank * sizeof (float))               /* c */
    );
}

static float*
gamma_pointer (uint8_t* base)
{
  Header* h = (Header*) base;

  return (float*) (
         base 
       + sizeof (*h) 
       + (h->n_raw_labels * sizeof (float))                      /* logpriorz */
       + (n_labels (h) * h->rank * sizeof (float))               /* c */
       + (n_labels (h) * h->rank * sizeof (float))               /* d */
    );
}

static const float*
logpriorz_pointer (const uint8_t* base)
{
  Header* h = (Header*) base;

  return (float*) (
         base 
       + sizeof (*h) 
    );
}

static size_t
file_size (const Header* h)
{
  return (
         sizeof (*h)                                             /* header */
       + (h->n_raw_labels * sizeof (float))                      /* logpriorz */
       + (n_labels (h) * h->rank * sizeof (float))               /* c */
       + (n_labels (h) * h->rank * sizeof (float))               /* d */
       + (h->rank * sizeof (float))                              /* gamma */
       + (h->n_workers * h->rank * sizeof (float))               /* loga */
    );
}

static float*
construct_logpriorz (const Header* h,
                     const float*  logpriorz)
{
  unsigned int n = n_labels (h);
  float* rv = malloc (n * sizeof (float));

  if (rv == NULL)
    {
      fprintf (stderr,
               "can't malloc %zu bytes: %s\n",
               n * sizeof (float),
               strerror (errno));
      return rv;
    }

  for (unsigned int i = 0; i < n; ++i)
    {
      unsigned int labels[h->n_raw_labels];
      unsigned int depth = index_to_set (i, labels, h->n_raw_labels);
      float logpz = 0.0;

      for (unsigned int j = 0; j < depth; ++j)
        { 
          logpz += logpriorz[labels[j]];
        }

      rv[i] = logpz;
    }

  safe_exponentiate_normalize (rv, n);

  for (unsigned int i = 0; i < n; ++i)
    {
      rv[i] = log (rv[i]);
    }

  return rv;
}


MultiModel*
multi_model_open (const char* pathname,
                  bool        read_only)
{
  Header tmp_hdr;
  int fd = open (pathname, (read_only) ? O_RDONLY : O_RDWR);

  if (fd < 0)
    {
      fprintf (stderr, "can't open '%s': %s\n", pathname, strerror (errno));
      return NULL;
    }

  if (lseek (fd, 0, SEEK_SET) == (off_t) -1 ||
      read (fd, &tmp_hdr, sizeof (tmp_hdr)) < (ssize_t) sizeof (tmp_hdr))
    {
      fprintf (stderr,
               "error reading header of model '%s': %s\n",
               pathname,
               strerror (errno));
      goto FAIL;
    }

  if (tmp_hdr.version != MODEL_VERSION)
    {
      fprintf (stderr, 
               "model version mismatch (want %llu got %llu)\n",
               MODEL_VERSION,
               (long long unsigned int) tmp_hdr.version);
      goto FAIL;
    }

  uint8_t* base = mmap (NULL,
                        file_size (&tmp_hdr),
                        PROT_READ | (read_only ? 0 : PROT_WRITE),
                        MAP_SHARED,
                        fd,
                        0);

  if (base == (uint8_t*) (-1))
    {
      fprintf (stderr,
               "can't mmap %zu bytes of '%s': %s\n",
               file_size (&tmp_hdr),
               pathname,
               strerror (errno));
      goto FAIL;
    }

  MultiModelImpl* impl = (MultiModelImpl*) malloc (sizeof (MultiModelImpl));
  if (impl == NULL)
    {
      fprintf (stderr,
               "wtf, can't alloc %zu bytes (?): %s\n",
               sizeof (MultiModelImpl),
               strerror (errno));
      goto FAIL;
    }
  impl->fd = fd;
  impl->base = base;
  impl->priorloga = parameterized_normal_distribution (1);
  impl->priorlogbeta = normal_distribution (0, 1);
  impl->priorc = normal_distribution (0, 1);
  impl->priord = normal_distribution (0, 1);

  MultiModel* multi_model = (MultiModel*) malloc (sizeof (MultiModel));
  if (multi_model == NULL)
    {
      fprintf (stderr,
               "wtf, can't alloc %zu bytes (?): %s\n",
               sizeof (MultiModel),
               strerror (errno));
      free (impl);
      goto FAIL;
    }

  impl->logpriorz = construct_logpriorz ((Header*) base,
                                         logpriorz_pointer (base));

  if (impl->logpriorz == NULL)
    {
      free (impl);
      goto FAIL;
    }

  multi_model->n_items = ((Header*) base)->n_items;
  multi_model->n_raw_labels = ((Header*) base)->n_raw_labels;
  multi_model->max_raw_labels = ((Header*) base)->max_raw_labels;
  multi_model->n_labels = n_labels ((Header*) base);
  multi_model->n_workers = ((Header*) base)->n_workers;
  multi_model->rank = ((Header*) base)->rank;
  multi_model->loga = loga_pointer (base);
  multi_model->gamma = gamma_pointer (base);
  multi_model->c = c_pointer (base);
  multi_model->d = d_pointer (base);
  multi_model->logpriorz = impl->logpriorz;
  multi_model->rawlogpriorz = logpriorz_pointer (base);
  multi_model->priorloga = &impl->priorloga;
  multi_model->priorlogbeta = &impl->priorlogbeta.base;
  multi_model->priorc = &impl->priorc.base;
  multi_model->priord = &impl->priord.base;
  multi_model->impl = impl;

  return multi_model;

FAIL:
  close (fd);
  return NULL;
}

MultiModel*
multi_model_creat (const char*        pathname,
                   unsigned int       n_items,
                   unsigned int       n_raw_labels,
                   unsigned int       max_raw_labels,
                   unsigned int       n_workers,
                   unsigned int       rank,
                   const float*       logpriorz)
{
  Header header = { .version = MODEL_VERSION,
                    .n_items = n_items,
                    .n_raw_labels = n_raw_labels,
                    .max_raw_labels = max_raw_labels,
                    .n_workers = n_workers,
                    .rank = rank };
  int fd = open (pathname,
                 O_RDWR | O_CREAT | O_EXCL,
                 S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH);

  if (fd < 0)
    {
      fprintf (stderr, "can't creat '%s': %s\n", pathname, strerror (errno));
      return NULL;
    }

  off_t end = lseek (fd, 0, SEEK_END);

  unsigned int n = n_labels (&header);
  float* init_c = (float*) malloc (n * rank * sizeof (float));
  float* init_d = (float*) malloc (n * rank * sizeof (float));

  for (unsigned int i = 0; i < n * rank; ++i)
    {
      init_c[i] = 0.1 * (drand48 () - 0.5);
      init_d[i] = 0.1 * (drand48 () - 0.5);
    }

  if (end != 0 ||
      ftruncate (fd, file_size (&header)) < 0 ||
      lseek (fd, 0, SEEK_SET) == (off_t) -1 ||
      write (fd, &header, sizeof (header)) < (ssize_t) sizeof (header) ||
      write (fd, logpriorz, n_raw_labels * sizeof (float)) < (ssize_t) (n_raw_labels * sizeof (float)) ||
      write (fd, init_c, n * rank * sizeof (float)) < (ssize_t) (n * rank * sizeof (float)) ||
      write (fd, init_d, n * rank * sizeof (float)) < (ssize_t) (n * rank * sizeof (float)))
    {
      fprintf (stderr,
               "error writing header of model '%s': %s\n",
               pathname,
               strerror (errno));
      unlink (pathname);
      close (fd);
      return NULL;
    }

  close (fd);
  return multi_model_open (pathname, false);
}

void
multi_model_close (MultiModel* multi_model)
{
  MultiModelImpl* impl = multi_model->impl;

  munmap (impl->base, file_size ((Header*) impl->base));
  close (impl->fd);
  free (impl);
  free (multi_model);
}
